#include <QMessageBox>

#include "../../mainwindow.h"
#include "../../datamanager/geometrydatanode.h"
#include "../../datamanager/proteindatanode.h"
#include "../../datamanager/reg3volumedatanode.h"
#include "../../sdf/pocketextractor.h"
#include "../../sdf/pocketrefine.h"
#include "../../contour/mcconextractor.h"
#include "../../renderer/rendermanager.h"
#include "../../data/fileplugin.h"
#include "pocketplugin.h"
#include "pocketcontrol.h"
#include "pocketrefinecontrol.h"

using namespace PBVP;

PocketPlugin::PocketPlugin()
{
	m_pControl = new PocketControl();
}

PocketPlugin::~PocketPlugin()
{
	delete m_pControl;
	delete progressBar;
}

void PocketPlugin::initialize(MainInterface* mainInterface)
{
	ToolPlugin::initialize(mainInterface);

	m_pDataManager = m_pMainInterface->getDataManager();
	m_pRenderManager = m_pMainInterface->getRenderManager();
	m_pFileExtManager = m_pMainInterface->getFileExtensionManager();

	pocketMenu = new QMenu(tr("Construct Pockets"), m_pMainInterface->getMainWindow());

	pocketSurfaceAction = new QAction(tr("Construct Pocket Envelops"), pocketMenu);
	pocketSurfaceAction->setStatusTip(description());
	connect(pocketSurfaceAction, SIGNAL(triggered()), this, SLOT(computePockets()));

	pocketVolumeAction = new QAction(tr("Construct Pocket Volume"), pocketMenu);
	pocketVolumeAction->setStatusTip(description());
	connect(pocketVolumeAction, SIGNAL(triggered()), this, SLOT(computePocketVolume()));

	pocketRefineAction = new QAction(tr("Refine a Pocket Envelop"), pocketMenu);
	pocketRefineAction->setStatusTip("Refine a pocket to better match the geometry of the protein surface");
	connect(pocketRefineAction, SIGNAL(triggered()), this, SLOT(refinePocket()));

	pocketMenu->addAction(pocketSurfaceAction);
	pocketMenu->addAction(pocketVolumeAction);
	pocketMenu->addAction(pocketRefineAction);

	progressBar = new QProgressBar();
	progressBar->setWindowTitle("Computing Pockets");
	progressBar->move(400, 400);
	progressBar->hide();
	progressBar->setMinimum(0);
	progressBar->setMaximum(100);
}

bool PocketPlugin::isSupported(DataNode* pNode)
{
#ifdef __GNU__
	if(pNode->data()->getDataSetType() == GEOMETRY_DATA) return true;
	if(pNode->data()->getDataSetType() == PROTEIN_DATA) return true;
#else
	// It supports either geometry data or protein data
	GeometryDataNode *pGeoNode = dynamic_cast<GeometryDataNode *>(pNode);
	if(pGeoNode) return true;
	ProteinDataNode *pProteinNode = dynamic_cast<ProteinDataNode *>(pNode);
	if(pProteinNode) return true;
#endif
	return false;
}

//////////////////////////////////////////////////////////////////////////
// public slots
//////////////////////////////////////////////////////////////////////////
void PocketPlugin::computePockets()
{
	DataNode *pNode = m_pDataManager->getCurrentNode();
	if (!pNode)
	{
		QMessageBox::critical(m_pMainInterface->getMainWindow(), "Error", "No Data Selected in DataManager");
		return;
	}
	if(!isSupported(pNode)) {
		QMessageBox::critical(m_pMainInterface->getMainWindow(), "Error", "Only can compute pockets for Protein Data or Geometry Data");
		return;
	}

	if (m_pControl->exec())
	{
		GeometryDataNode *pGeoNode = NULL;
#ifdef __GNU__
		if(pNode->data()->getDataSetType() == GEOMETRY_DATA) pGeoNode = (GeometryDataNode *)(pNode);
#else
		pGeoNode = dynamic_cast<GeometryDataNode *>(pNode);
#endif
		if (pGeoNode) {	// It is a Geometry Node
			progressBar->setValue(0);
			progressBar->show();
			GeometryData *pGeoData = (GeometryData *)(pGeoNode->data());
			PocketExtractor *pPocketExtract = new PocketExtractor(m_pControl->getMarchDist(),
				m_pControl->getIsoVal(), m_pControl->getFilterIterations());
			pPocketExtract->setProgressBar(progressBar);
			GeometryData *pPocketSurf = pPocketExtract->computePocketSurface(pGeoData);
			delete pPocketExtract;
			progressBar->hide();

			GeometryDataNode *pGeoNode = new GeometryDataNode(pPocketSurf, m_pMainInterface);
			pGeoNode->setName(pNode->getName().append("_pocket"));
			//m_pRenderManager->setDefaultRenderer(pGeoNode);
			m_pDataManager->addDataNode(pGeoNode);
			return;
		}

		ProteinDataNode *pProteinNode = NULL;
#ifdef __GNU__
		if(pNode->data()->getDataSetType() == PROTEIN_DATA) pProteinNode = (ProteinDataNode *)(pNode);
#else
		pProteinNode = dynamic_cast<ProteinDataNode *>(pNode);
#endif
		if(pProteinNode) {	// It is a Protein Node
			progressBar->setValue(0);
			progressBar->show();
			ProteinData *pProtein = (ProteinData *)(pProteinNode->data());
			PocketExtractor *pPocketExtract = new PocketExtractor(m_pControl->getMarchDist(),
				m_pControl->getIsoVal(), m_pControl->getFilterIterations());
			pPocketExtract->setProgressBar(progressBar);
			GeometryData *pPocketSurf = pPocketExtract->computePocketSurface(pProtein->getAtoms());
			delete pPocketExtract;
			progressBar->hide();

			GeometryDataNode *pGeoNode = new GeometryDataNode(pPocketSurf, m_pMainInterface);
			pGeoNode->setName(pNode->getName().append("_pocket"));
			//m_pRenderManager->setDefaultRenderer(pGeoNode);
			m_pDataManager->addDataNode(pGeoNode);
			return;
		}
	}
}

void PocketPlugin::computePocketVolume()
{
	DataNode *pNode = m_pDataManager->getCurrentNode();
	if (!pNode)
	{
		QMessageBox::critical(m_pMainInterface->getMainWindow(), "Error", "No Data Selected in DataManager");
		return;
	}
	if(!isSupported(pNode)) {
		QMessageBox::critical(m_pMainInterface->getMainWindow(), "Error", "Only can compute pockets for Protein Data or Geometry Data");
		return;
	}

	Reg3DVolumeDataNode* pVolNode;
	Reg3DVolumeData* pPocketVol;

	if (m_pControl->exec())
	{
		progressBar->setValue(0);
		progressBar->show();
		GeometryDataNode *pGeoNode = NULL;
		#ifdef __GNU__
				if(pNode->data()->getDataSetType() == GEOMETRY_DATA) pGeoNode = (GeometryDataNode *)(pNode);
		#else
				pGeoNode = dynamic_cast<GeometryDataNode *>(pNode);
		#endif
		if (pGeoNode) {	// It is a Geometry Node
			GeometryData *pGeoData = (GeometryData *)(pGeoNode->data());
			PocketExtractor *pPocketExtract = new PocketExtractor(m_pControl->getMarchDist(),
				m_pControl->getIsoVal(), m_pControl->getFilterIterations());
			pPocketExtract->setProgressBar(progressBar);
			pPocketVol = pPocketExtract->computePocketVolume(pGeoData);
			delete pPocketExtract;
		}

		ProteinDataNode *pProteinNode = NULL;
		#ifdef __GNU__
				if(pNode->data()->getDataSetType() == PROTEIN_DATA) pProteinNode = (ProteinDataNode *)(pNode);
		#else
				pProteinNode = dynamic_cast<ProteinDataNode *>(pNode);
		#endif
		if(pProteinNode) {	// It is a Protein Node
			ProteinData *pProtein = (ProteinData *)(pProteinNode->data());
			PocketExtractor *pPocketExtract = new PocketExtractor(m_pControl->getMarchDist(),
				m_pControl->getIsoVal(), m_pControl->getFilterIterations());
			pPocketExtract->setProgressBar(progressBar);
			pPocketVol = pPocketExtract->computePocketVolume(pProtein->getAtoms());
			delete pPocketExtract;
		}
		progressBar->hide();

		pVolNode = new Reg3DVolumeDataNode(pPocketVol, m_pMainInterface);
		pVolNode->setName(pNode->getName().append("_pocket.rawiv"));
		m_pRenderManager->setDefaultRenderer(pVolNode);
		m_pDataManager->addDataNode(pVolNode);
	}
}

void PocketPlugin::refinePocket()
{
	PocketRefineControl *refineControl = new PocketRefineControl(m_pFileExtManager);
	if(refineControl->exec()) {
		printf("pocket file: %s, protein file: %s, epsilon = %f\n", (refineControl->getPocketFileName()).toAscii().data(),
			(refineControl->getProteinFileName()).toAscii().data(), refineControl->getEpsilon());
		QString pocketFname = refineControl->getPocketFileName();
		QString proteinFname = refineControl->getProteinFileName();
		FilePluginInterface *pFilePlugin = m_pFileExtManager->getSupportedPlugin(pocketFname);
		if(!pFilePlugin) {
			delete refineControl;
			return;
		}
		GeometryData *pPocket = (GeometryData *)pFilePlugin->loadFile(pocketFname);
		if(!pPocket) {
			delete refineControl;
			return;
		}
		pFilePlugin = m_pFileExtManager->getSupportedPlugin(proteinFname);
		if(!pFilePlugin) {
			delete pPocket;
			delete refineControl;
			return;
		}
		GeometryData *pProtein = (GeometryData *)pFilePlugin->loadFile(proteinFname);
		if(!pProtein) {
			delete pPocket;
			delete refineControl;
			return;
		}
		// compute the volume pocket function, which is needed for refinement
		PocketExtractor *pPocketExtract = new PocketExtractor(m_pControl->getMarchDist(),
			m_pControl->getIsoVal(), m_pControl->getFilterIterations());
		Reg3DVolumeData* pPocketVol = pPocketExtract->computePocketVolume(pProtein);
		delete pPocketExtract;

		PocketRefine refine(pProtein, pPocketVol);
		GeometryData *surf = refine.refinePocket(pPocket, refineControl->getEpsilon());
		if (surf)
		{
			GeometryDataNode *pGeoNode = new GeometryDataNode(surf, m_pMainInterface);
			QString fname = pocketFname.append("_refined");
			pGeoNode->setName(fname);
			m_pDataManager->addDataNode(pGeoNode);
		}
		delete pPocketVol;
		delete pProtein;
		delete pPocket;
	}
	delete refineControl;
}

Q_EXPORT_PLUGIN2(pocket, PocketPlugin)
